#!/usr/bin/env python
import rospy
# import keyboard

from geometry_msgs.msg import Twist

from webots_ros.srv import set_float
from webots_ros.srv import get_float


webots_ns = 'sdnyc_webots'
TIMESTEP = 32

# lfwp = [-0.501372, -0.399851, -0.202694]
# rfwp = [-0.502124, 0.399914, -0.197207]
# lbwp = [0.499075, -0.400031, -0.202388]
# rbwp = [0.498591, 0.399956, -0.196755]
W = 0.394 # wheel width
R = 0.15

class MotorManager():
    def __init__(self, is_persistent = False):
        f_l_wheel_client = '{}/front_left_wheel/{}'
        f_r_wheel_client = '{}/front_right_wheel/{}'
        b_l_wheel_client = '{}/back_left_wheel/{}'
        b_r_wheel_client = '{}/back_right_wheel/{}'
        
        # create proxy methods to set motor positions
        rospy.wait_for_service(f_l_wheel_client.format(webots_ns, 'set_position'))
        self.fl_set_position = rospy.ServiceProxy(
            f_l_wheel_client.format(webots_ns, 'set_position'), set_float, persistent=is_persistent)
        rospy.wait_for_service(f_r_wheel_client.format(webots_ns, 'set_position'))
        self.fr_set_position = rospy.ServiceProxy(
            f_r_wheel_client.format(webots_ns, 'set_position'), set_float, persistent=is_persistent)
        rospy.wait_for_service(b_l_wheel_client.format(webots_ns, 'set_position'))
        self.bl_set_position = rospy.ServiceProxy(
            b_l_wheel_client.format(webots_ns, 'set_position'), set_float, persistent=is_persistent)
        rospy.wait_for_service(b_r_wheel_client.format(webots_ns, 'set_position'))
        self.br_set_position = rospy.ServiceProxy(
            b_r_wheel_client.format(webots_ns, 'set_position'), set_float, persistent=is_persistent)

        # create proxy methods to set motor velocities
        rospy.wait_for_service(f_l_wheel_client.format(webots_ns, 'set_velocity'))
        self.fl_set_velocity = rospy.ServiceProxy(
            f_l_wheel_client.format(webots_ns, 'set_velocity'), set_float, persistent=is_persistent)
        rospy.wait_for_service(f_r_wheel_client.format(webots_ns, 'set_velocity'))
        self.fr_set_velocity = rospy.ServiceProxy(
            f_r_wheel_client.format(webots_ns, 'set_velocity'), set_float, persistent=is_persistent)
        rospy.wait_for_service(b_l_wheel_client.format(webots_ns, 'set_velocity'))
        self.bl_set_velocity = rospy.ServiceProxy(
            b_l_wheel_client.format(webots_ns, 'set_velocity'), set_float, persistent=is_persistent)
        rospy.wait_for_service(b_r_wheel_client.format(webots_ns, 'set_velocity'))
        self.br_set_velocity = rospy.ServiceProxy(
            b_r_wheel_client.format(webots_ns, 'set_velocity'), set_float, persistent=is_persistent)


        # create proxy methods to get motor velocities
        rospy.wait_for_service(f_l_wheel_client.format(webots_ns, 'get_velocity'))
        self.fl_get_velocity = rospy.ServiceProxy(
            f_l_wheel_client.format(webots_ns, 'get_velocity'), get_float, persistent=is_persistent)
        rospy.wait_for_service(f_r_wheel_client.format(webots_ns, 'get_velocity'))
        self.fr_get_velocity = rospy.ServiceProxy(
            f_r_wheel_client.format(webots_ns, 'get_velocity'), get_float, persistent=is_persistent)
        rospy.wait_for_service(b_l_wheel_client.format(webots_ns, 'get_velocity'))
        self.bl_get_velocity = rospy.ServiceProxy(
            b_l_wheel_client.format(webots_ns, 'get_velocity'), get_float, persistent=is_persistent)
        rospy.wait_for_service(b_r_wheel_client.format(webots_ns, 'get_velocity'))
        self.br_get_velocity = rospy.ServiceProxy(
            b_r_wheel_client.format(webots_ns, 'get_velocity'), get_float, persistent=is_persistent)

        self.fl_set_position(float('inf'))
        self.bl_set_position(float('inf'))

        self.fr_set_position(float('inf'))
        self.br_set_position(float('inf'))

        self.__keyboard_topic = '/key_vel'
        rospy.Subscriber(self.__keyboard_topic, Twist, callback=self.__keyboard_callback, queue_size=1, buff_size=1)

        self.set_once = False

        rospy.loginfo('started.')

    def __keyboard_callback(self, msg):
        v = msg.linear.x
        w = msg.angular.z

        if v != 0 or w !=0:
            rospy.loginfo('v={} | w={}. | msg={}'.format(v, w, msg))
            self.set_once = True
            cm.set_vw(v=v, w=w)
        else: # v==w==0
            if self.set_once:
                self.set_once = False
                cm.set_vw(v=v, w=w)

    def set_velocity(self, l_speed, r_speed):

        # if (l_speed >= 0.0):
        #     self.fl_set_position(float('inf'))
        #     self.bl_set_position(float('inf'))
        # else:
        #     self.fl_set_position(float('-inf'))
        #     self.bl_set_position(float('-inf'))
        
        # if (r_speed >= 0.0):
        #     self.fr_set_position(float('inf'))
        #     self.br_set_position(float('inf'))
        # else:
        #     self.fr_set_position(float('-inf'))
        #     self.br_set_position(float('-inf'))
        
        self.fl_set_velocity(l_speed)
        self.bl_set_velocity(l_speed)

        self.fr_set_velocity(r_speed)
        self.br_set_velocity(r_speed)

        # rospy.loginfo('vel set!')

    def set_vw(self, v, w):
        u_r = (2*v + W*w)/(2*R)
        u_l = (2*v - W*w)/(2*R)

        self.set_velocity(l_speed = u_l, r_speed = u_r)
        # print(f'u_l:{u_l} | u_r:{u_r}')

    def get_vw(self):
        # u_l = self.l_avg_spd
        # u_r = self.r_avg_spd
        u_l = self.fl_get_velocity(0).value
        u_r = self.fr_get_velocity(0).value
        # print(f'u_l:{u_l} | u_r:{u_r}')

        w = R * (u_r - u_l) / W
        v = R * (u_r + u_l) / 2

        return (v, w)

        
if __name__ == "__main__":

    rospy.init_node('MotorManager', anonymous= False)

    cm = MotorManager()

    rospy.spin()
    # while not rospy.is_shutdown():
    #     if keyboard.is_pressed('w'):
    #         cm.set_vw(v=0.8, w=0)
    #     if keyboard.is_pressed('s'):
    #         cm.set_vw(v=-0.5, w=0)
    #     if keyboard.is_pressed('a'):
    #         cm.set_vw(v=0, w=1.0)
    #     if keyboard.is_pressed('d'):
    #         cm.set_vw(v=0, w=-1.0)

